// MyNpzScoreWriter.cc
#include "MyNpzScoreWriter.hh"

#include <vector>

#include "G4SystemOfUnits.hh"
#include "G4VScoringMesh.hh"

void MyNpzScoreWriter::DumpQuantityToFile(const G4String& psName, const G4String& fileName, const G4String&)
{
    if (!fScoringMesh) return;

    // 1. 拿到 map
    auto scoreMap = fScoringMesh->GetScoreMap();
    auto it = scoreMap.find(psName);
    if (it == scoreMap.end()) return;
    // auto score = it->second->GetMap();
    std::map<G4int, G4StatDouble*>* score = it->second->GetMap();

    // 2. 建一个 3D 数组，顺序和 Geant4 内部一致
    std::vector<double> buf(fNMeshSegments[0] * fNMeshSegments[1] * fNMeshSegments[2], 0.0);

    G4double unitValue = fScoringMesh->GetPSUnitValue(psName);
    for (int ix = 0; ix < fNMeshSegments[0]; ++ix)
    {
        for (int iy = 0; iy < fNMeshSegments[1]; ++iy)
        {
            for (int iz = 0; iz < fNMeshSegments[2]; ++iz)
            {
                G4int idx = GetIndex(ix, iy, iz);  // 父类提供
                auto v = score->find(idx);
                auto ss = *(v->second);
                if (v != score->end()) buf[idx] = *(v->second) / unitValue;
            }
        }
    }

    // 3. 写成 npy，shape = (Nx,Ny,Nz)
    cnpy::npy_save(fileName, buf.data(), {size_t(fNMeshSegments[0]), size_t(fNMeshSegments[1]), size_t(fNMeshSegments[2])}, "w");
}